require("stack")

GRAPH = {}
local GraphMeta = {
	__index = GRAPH
}
local oldPrint = print
local function print( ... )
	--oldPrint( ... )
end


function GRAPH.New()
	local t = {
		Nodes = {},
		Edges = {}
	}
	
	setmetatable( t, GraphMeta )

	return t
end

function GRAPH:GetNode( name, throwInexistantError )
	
	if self.Nodes[ name ] == nil then
		if throwInexistantError == true then
			error("Node " .. name .. " does not exist")
		end
	
		self.Nodes[ name ] = NODE.New( name )
	end
	
	return self.Nodes[ name ]
	
end

function GRAPH:AddEdge( edge )
	
	local source = self:GetNode( edge.src )
	local destination = self:GetNode( edge.dst )
	local weight = edge.miles
	
	local newEdge = EDGE.New( source, destination, weight )
	
	table.insert( self.Edges, newEdge )
	
	source:AddEdge( newEdge )
	destination:AddBackwardEdge( newEdge )
	
end

function GRAPH:Reset()
	for _, edge in pairs( self.Edges ) do
		edge:Reset()
	end
end

function GRAPH:Solidify( Stack, maximunFlow )

	print("\t\tSolidifying maximunFlow of: " .. maximunFlow )
	print("\t\t" .. tostring(Stack) )
	
	for i=1, Stack.top do
		local node = Stack.list[ i ]
		local edge = node.ActiveEdge
		
		if edge.Source == node then
			edge:AddFlow( maximunFlow )
			print("\t\t\tForward: " .. tostring(node) )
		else
			edge:AddFlow( -maximunFlow ) --If we are going backwards, we have to subtract
			print("\t\t\tBackward: " .. tostring(node) )
		end
	
	end
	
end

function GRAPH:MaxFlow( sourceName, destinationName )
	
	local source = self:GetNode( sourceName, true )
	local destination = self:GetNode( destinationName, true )
	local Stack = stack.New()
	
	local function Explore( node, maximunFlow )
	
		print("Exploring node: " .. tostring(node) .. " (Flow: ".. tostring(maximunFlow) .. ")")
		
		if maximunFlow and maximunFlow < 0 then
			error("Invalid maximunFlow.")
		elseif maximunFlow == 0 then
			return 0
		end
	
		if node == destination then --We have found the pot of gold!
		
			--Go back through the whole stack and add it to the flow of the graph
			self:Solidify( Stack, maximunFlow )
			
			return maximunFlow
		end
		
		if Stack:has( node ) then --Oh uh.. We have a cycle, let's skip this.
			print("Cycle found")
			
			return 0
		end
		
		--Push the node into the stack, so we know our steps
		Stack:push( node )
		
		local SumUsed = 0 --Total amount we are going to send through
	
		for _, edge in ipairs( node.Edges ) do
		
			node.ActiveEdge = edge
			local residual = edge:GetResidual( node ) --The amount that the tube can still hold
			local flow = maximunFlow and math.min( maximunFlow, residual ) or residual --The amount that we are going to transfer through it
			
			print( "\t Looking at edge: " .. tostring(edge)  )
			
			if flow > 0 then
				local used = Explore( edge.Target, flow ) --Amount that it is going to go through the tube
				
				if maximunFlow then
					maximunFlow = maximunFlow - used --Recalculate the amount we still have left
				end
				
				SumUsed = SumUsed + used
			end
		
		end
		
		--Remove the node from the stack, to erase our steps
		Stack:pop()
		
		return SumUsed
	
	end
	
	Explore( source )
	
	local Result = {}
	
	for _, edge in pairs( self.Edges ) do
		Result[ edge ] = edge.Flow
	end
	
	self:Reset()
	
	return Result

end


NODE = {}
local NodeMeta = {
	__index = NODE,
	__tostring = function(node) return node.Name end
}

function NODE.New( name )
	local t = {
		Name = name,
		Edges = {},
		BackwardEdges = {}
	}
	
	setmetatable( t, NodeMeta )

	return t
end

function NODE:AddEdge( edge )
	table.insert( self.Edges, edge )
end


function NODE:AddBackwardEdge( edge )
	table.insert( self.Edges, edge )
end

EDGE = {}
local EdgeMeta = {
	__index = EDGE,
	__tostring = function(edge)
		return "EDGE(" .. edge.Flow .. "/" .. edge.Capacity .. " - " .. tostring(edge.Source) .. "->" .. tostring(edge.Target) .. ")"
	end
}

function EDGE.New( source, target, total )
	local t = {
		Source = source,
		Target = target,
		Capacity = total,
		Flow = 0
	}
	
	setmetatable( t, EdgeMeta )

	return t	
end

function EDGE:AddFlow( flow )
	self.Flow = self.Flow + flow
	
	if self.Flow > self.Capacity then
		error("Node passaing maximun capacity!")
	elseif self.Flow < 0 then
		error("Oh no! A black hole! We can not have negative matter!")
	end
end

function EDGE:GetResidual( node )
	if node == self.Source then
		return self.Capacity - self.Flow
	else --We are going backwards! The residual is we are actually using
		return self.Flow
	end
end

function EDGE:Reset()
	self.Flow = 0
end
